#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>

// 矩阵维度配置（可修改）
#define ROWS 1024  // 矩阵行数
#define COLS 1024  // 矩阵列数

// CUDA错误检查宏
#define CHECK(call) {                                                         \
    const cudaError_t error = call;                                            \
    if (error != cudaSuccess) {                                                \
        printf("Error: %s:%d, ", __FILE__, __LINE__);                         \
        printf("code:%d, reason: %s\n", error, cudaGetErrorString(error));    \
        exit(1);                                                               \
    }                                                                          \
}

// GPU核函数：矩阵相加
__global__ void matrixAddKernel(float* A, float* B, float* C, int rows, int cols) {
    // 计算全局行列索引
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    // 边界检查
    if (row < rows && col < cols) {
        int index = row * cols + col;
        C[index] = A[index] + B[index];
    }
}

// CPU参考实现：矩阵相加（用于验证）
void matrixAddCPU(float* A, float* B, float* C, int rows, int cols) {
    for (int row = 0; row < rows; ++row) {
        for (int col = 0; col < cols; ++col) {
            int index = row * cols + col;
            C[index] = A[index] + B[index];
        }
    }
}

// 初始化矩阵（随机值）
void initMatrix(float* mat, int size) {
    for (int i = 0; i < size; ++i) {
        mat[i] = static_cast<float>(rand()) / RAND_MAX; // [0, 1]范围随机数
    }
}

// 验证结果（比较GPU和CPU计算结果）
bool verifyResult(float* C_ref, float* C_gpu, int size) {
    const float epsilon = 1e-5f;
    for (int i = 0; i < size; ++i) {
        if (fabs(C_ref[i] - C_gpu[i]) > epsilon) {
            printf("Mismatch at index %d: CPU=%.5f, GPU=%.5f\n", 
                  i, C_ref[i], C_gpu[i]);
            return false;
        }
    }
    return true;
}

int main() {
    // 矩阵元素总数
    const int matrixSize = ROWS * COLS;

    // 分配主机内存
    float *h_A = (float*)malloc(matrixSize * sizeof(float));
    float *h_B = (float*)malloc(matrixSize * sizeof(float));
    float *h_C = (float*)malloc(matrixSize * sizeof(float));
    float *h_C_ref = (float*)malloc(matrixSize * sizeof(float));

    // 初始化输入矩阵
    srand(2023);
    initMatrix(h_A, matrixSize);
    initMatrix(h_B, matrixSize);

    // 分配设备内存
    float *d_A, *d_B, *d_C;
    CHECK(cudaMalloc(&d_A, matrixSize * sizeof(float)));
    CHECK(cudaMalloc(&d_B, matrixSize * sizeof(float)));
    CHECK(cudaMalloc(&d_C, matrixSize * sizeof(float)));

    // 拷贝数据到设备
    CHECK(cudaMemcpy(d_A, h_A, matrixSize * sizeof(float), cudaMemcpyHostToDevice));
    CHECK(cudaMemcpy(d_B, h_B, matrixSize * sizeof(float), cudaMemcpyHostToDevice));

    // 设置执行配置
    dim3 block(32, 32);  // 每个块包含 1024 个线程 (32x32)
    dim3 grid((COLS + block.x - 1) / block.x, 
             (ROWS + block.y - 1) / block.y);

    // 执行核函数并计时
    cudaEvent_t start, stop;
    CHECK(cudaEventCreate(&start));
    CHECK(cudaEventCreate(&stop));
    
    CHECK(cudaEventRecord(start));
    matrixAddKernel<<<grid, block>>>(d_A, d_B, d_C, ROWS, COLS);
    CHECK(cudaEventRecord(stop));
    CHECK(cudaEventSynchronize(stop));
    
    float milliseconds = 0;
    CHECK(cudaEventElapsedTime(&milliseconds, start, stop));
    printf("GPU Time: %.3f ms\n", milliseconds);

    // 拷贝结果回主机
    CHECK(cudaMemcpy(h_C, d_C, matrixSize * sizeof(float), cudaMemcpyDeviceToHost));

    // CPU计算参考结果
    matrixAddCPU(h_A, h_B, h_C_ref, ROWS, COLS);

    // 验证结果
    if (verifyResult(h_C_ref, h_C, matrixSize)) {
        printf("Result Verification: PASSED\n");
    }

    // 释放资源
    free(h_A);
    free(h_B);
    free(h_C);
    free(h_C_ref);
    CHECK(cudaFree(d_A));
    CHECK(cudaFree(d_B));
    CHECK(cudaFree(d_C));
    CHECK(cudaEventDestroy(start));
    CHECK(cudaEventDestroy(stop));

    return 0;
}
